Skip to content

refactor: remove coremltools dependency and use torch.utils.cpp_extension to load kmeans1d#31

Merged
guru-desh merged 19 commits into
apple:mainfrom
guru-desh:u/gurudesh/remove-coremltools-dependency-torch-cpp
Jul 15, 2026
Merged

refactor: remove coremltools dependency and use torch.utils.cpp_extension to load kmeans1d#31
guru-desh merged 19 commits into
apple:mainfrom
guru-desh:u/gurudesh/remove-coremltools-dependency-torch-cpp

Conversation

@guru-desh

@guru-desh guru-desh commented Jul 9, 2026

Copy link
Copy Markdown
Contributor

Description

This PR removes coremltools as a required dependency for coreai-opt.

The only part of coremltools that was being used was kmeans1d. We have copied kmeans1d in coreai-opt using torch.utils.cpp_extension to compile the C++ code upon invocation.

Using torch.utils.cpp_extension allows us to have a more faithful refactor for kmeans1d. From here, we can then choose to explore alternatives like numba or optimizations for the kmeans1d C++ code. This PR is scoped to have the minimal amount of changes needed to remove coremltools as a dependency.

Thanks @pkmandke for the discussion and suggestions!

Testing

  1. We have equivalence tests to ensure that this copy of kmeans1d is equivalent to coremltools kmeans1d (tolerance is 1e-9)
  2. kmeans1d tests are copied over from coremltools
  3. We also have a hash check to check that the C++ file copied over is the same as the one that exists on coremltools. This means that we have a pure refactor of kmeans1d

@guru-desh
guru-desh force-pushed the u/gurudesh/remove-coremltools-dependency-torch-cpp branch from cca4b22 to 9910570 Compare July 9, 2026 21:03
_DLL = None


def _dll():

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This file is mainly the same as coremltools._deps.kmeans1d except for this function that uses torch.utils.cpp_extension to compile _core.cpp

Comment thread src/coreai_opt/deps/_kmeans1d/core.py
@guru-desh
guru-desh force-pushed the u/gurudesh/remove-coremltools-dependency-torch-cpp branch 3 times, most recently from 8c4d108 to a17f5c0 Compare July 10, 2026 00:05
# NOT inherit sysconfig flags and defaults to -O0, which made this
# tight-inner-loop, template-heavy DP code 10-15x slower per call than
# coremltools' prebuilt extension. Pass the same flags explicitly to match.
_EXTRA_CFLAGS = ["-std=c++11", "-O2", "-DNDEBUG"] + (

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Future PRs can look into using -O3 compilation flag and benchmark performance, but since this PR is meant as a pure refactor, we stick with -O2

@guru-desh
guru-desh force-pushed the u/gurudesh/remove-coremltools-dependency-torch-cpp branch from e858d3a to a64cf0c Compare July 10, 2026 09:08
@guru-desh guru-desh added the enhancement New feature or request label Jul 10, 2026
Comment thread src/coreai_opt/deps/_kmeans1d/core.py
@guru-desh
guru-desh force-pushed the u/gurudesh/remove-coremltools-dependency-torch-cpp branch from 1d5cdaf to a14c10c Compare July 13, 2026 00:18
@guru-desh
guru-desh marked this pull request as ready for review July 13, 2026 20:05
Comment thread changelog.d/31.changed Outdated
@guru-desh
guru-desh force-pushed the u/gurudesh/remove-coremltools-dependency-torch-cpp branch from 74244d2 to de2bd23 Compare July 14, 2026 06:09
Comment thread changelog.d/31.changed Outdated
Comment thread src/coreai_opt/deps/_kmeans1d/core.py Outdated
Comment thread tests/palettization/test_kmeans1d.py Outdated
Comment thread pyproject.toml Outdated
Comment thread pyproject.toml
Comment thread tests/palettization/test_kmeans1d.py Outdated
Comment thread tests/palettization/test_kmeans1d.py Outdated
guru-desh added a commit to guru-desh/coreai-optimization that referenced this pull request Jul 14, 2026
Drop "(copied byte-for-byte from coremltools)" -- attribution already
lives in the vendored files' headers and NOTICE.txt, and the phrase
will go stale once the core is further optimized (per aseemw's PR apple#31
review).
guru-desh added a commit to guru-desh/coreai-optimization that referenced this pull request Jul 14, 2026
aseemw's PR apple#31 review: kmeans1d correctness tests importing
coremltools as a live oracle defeats the purpose of removing it as a
required dependency, and the "must stay byte-identical to upstream"
canary asserts an invariant that isn't actually required (the vendored
code can be modified/adapted going forward).

Collapses TestEquivalence/TestKmeans1DBehavior/TestsKmeans1dFromCoremltools
into one TestKmeans1D class and deletes TestVendoredSourceMatchesUpstream
entirely. Small discrete-input cases (k>n clamping, fewer-than-k-cluster
collapse, explicit duplicates) now assert hardcoded literal values
directly, captured by running the vendored cluster() once (already
validated against coremltools in CI, so a known-good golden value).

Larger/random-array cases (unweighted/weighted equivalence grids,
large-n up to 100k, fp16 duplicates, weighted-integer-counts) would
require thousands of literal values to hardcode inline, so their
expected coremltools-oracle output is instead checked into
tests/palettization/assets/kmeans1d/*.npz (generated once via the new
generate_assets.py, not collected by pytest) and loaded at test time --
inputs are regenerated from the same fixed RNG seeds rather than also
stored on disk, keeping the assets small (~300KB total).

Net effect: this test file has zero coremltools imports (confirmed via
AST inspection), so coreml stops being an implicit test dependency for
kmeans1d correctness.
guru-desh added a commit to guru-desh/coreai-optimization that referenced this pull request Jul 14, 2026
aseemw's PR apple#31 review: kmeans1d correctness tests importing
coremltools as a live oracle defeats the purpose of removing it as a
required dependency, and the "must stay byte-identical to upstream"
canary asserts an invariant that isn't actually required (the vendored
code can be modified/adapted going forward).

Collapses TestEquivalence/TestKmeans1DBehavior/TestsKmeans1dFromCoremltools
into one TestKmeans1D class and deletes TestVendoredSourceMatchesUpstream
entirely. Small discrete-input cases (k>n clamping, fewer-than-k-cluster
collapse, explicit duplicates) now assert hardcoded literal values
directly, captured by running the vendored cluster() once (already
validated against coremltools in CI, so a known-good golden value).

Larger/random-array cases (unweighted/weighted equivalence grids,
large-n up to 100k, fp16 duplicates, weighted-integer-counts) would
require thousands of literal values to hardcode inline, so their
expected coremltools-oracle output is instead checked into
tests/palettization/assets/kmeans1d/*.npz and loaded at test time --
inputs are regenerated from the same fixed RNG seeds rather than also
stored on disk, keeping the assets small (~300KB total).

Net effect: this test file has zero coremltools imports (confirmed via
AST inspection), so coreml stops being an implicit test dependency for
kmeans1d correctness.
@guru-desh
guru-desh force-pushed the u/gurudesh/remove-coremltools-dependency-torch-cpp branch from 39c0fb8 to 9541614 Compare July 14, 2026 22:49
guru-desh added a commit to guru-desh/coreai-optimization that referenced this pull request Jul 14, 2026
aseemw's PR apple#31 review: kmeans1d correctness tests importing
coremltools as a live oracle defeats the purpose of removing it as a
required dependency, and the "must stay byte-identical to upstream"
canary asserts an invariant that isn't actually required (the vendored
code can be modified/adapted going forward).

Collapses TestEquivalence/TestKmeans1DBehavior/TestsKmeans1dFromCoremltools
into one TestKmeans1D class and deletes TestVendoredSourceMatchesUpstream
entirely. Small discrete-input cases (k>n clamping, fewer-than-k-cluster
collapse, explicit duplicates) now assert hardcoded literal values
directly, captured by running the vendored cluster() once (already
validated against coremltools in CI, so a known-good golden value).

Larger/random-array cases (unweighted/weighted equivalence grids,
large-n up to 100k, fp16 duplicates, weighted-integer-counts) would
require thousands of literal values to hardcode inline, so their
inputs and expected coremltools-oracle output are instead checked into
tests/palettization/assets/kmeans1d/*.npz and loaded at test time
(~3.3MB total, mostly the large-n case).

Net effect: this test file has zero coremltools imports (confirmed via
AST inspection), so coreml stops being an implicit test dependency for
kmeans1d correctness.
@guru-desh
guru-desh force-pushed the u/gurudesh/remove-coremltools-dependency-torch-cpp branch from 9541614 to 416db85 Compare July 14, 2026 22:54
Vendor coremltools' kmeans1d C++ core (_core.cpp, byte-for-byte from apple/coremltools; upstream kmeans1d 0.3.1, MIT) and JIT-compile it at runtime with torch.utils.cpp_extension, invoking its extern "C" entry points via ctypes. Provides cluster()/Clustered as a drop-in replacement for coremltools._deps._kmeans1d, with fail-loud input validation.
guru-desh added 18 commits July 15, 2026 11:24
Point _efficient_kmeans and kmeans_fake_palettize at coreai_opt._utils._kmeans1d instead of coremltools._deps._kmeans1d. The cluster()/Clustered call contract is unchanged.
Drop the optional third-party kmeans1d import and _HAS_KMEANS1D flag in palettize_utils; import coreai_opt._utils._kmeans1d unconditionally. The sklearn fallback for vector k-means is unchanged.
Compare cluster() output against the coremltools oracle (unweighted, weighted, large-n, collapse, fp16 ties), assert fixed known-good cases and input validation, and add a canary that hard-fails if the vendored _core.cpp drifts from apple/coremltools main.
coremltools is no longer needed at runtime; move it (and its numpy<2.4 cap) into a coreml optional-dependency + self-referencing dependency-group, kept in default-groups so the test oracle is available. Ship the vendored _core.cpp and LICENSE as package data.
The vendored kmeans1d C++ extension was 10-15x slower per call than
coremltools' precompiled version (measured directly: ~0.9s vs ~0.06s
on a realistic-size weighted cluster() call), even with the actual
JIT-compile lock/cache fully warmed. This was the real cause of the
1.2x-2.6x palettization wall-clock regression seen in an end-to-end
LLM eval, not JIT-compile lock contention across parallel workers
(a prior "warmup()" fix targeting that, since reverted, made no
measurable difference).

Root cause: torch.utils.cpp_extension.load() does not inherit
CPython's sysconfig OPT/CFLAGS the way distutils/setuptools does, so
it defaults to -O0. coremltools' setup.py sets no explicit -O flag
either, but distutils prepends sysconfig's "-O2 -DNDEBUG" (the
standard python.org/conda default) ahead of its extra_compile_args,
so its shipped .so is effectively an -O2 build. This tight-inner-loop,
template-heavy DP algorithm is extremely sensitive to that gap.

Pass the same flags explicitly so torch's JIT build matches
coremltools' effective compile line. Numeric output is unaffected
(verified against the existing coremltools-oracle equivalence tests).
torch.utils.cpp_extension.load() imports setuptools internally to JIT-compile
the vendored kmeans1d core. This was previously pulled in transitively via
coremltools, which is now an optional dependency, so it must be declared
explicitly.
Isolates the vendored third-party kmeans1d source (and its MIT LICENSE)
into a dedicated deps/ folder, separate from coreai-opt's own internal
_utils modules, per Apple legal's guidance on incorporating vendored
dependencies (mirrors how coremltools itself vendored the same code).
Completes the previous move: repoints pyproject.toml package-data and
the three call sites at coreai_opt.deps._kmeans1d instead of the old
coreai_opt._utils._kmeans1d location. Also anchors the .gitignore
deps/ rule to the repo root (/deps/) so it stops shadowing the new
src/coreai_opt/deps/ package.
Bring the per-file headers in line with Apple legal's guidance for
incorporating vendored third-party code (mirrors coremltools' own
kmeans1d headers): every file in the vendored package now carries the
upstream MIT license text followed by a trailing Apple copyright line,
with no coreai-opt repo-wide header layered on top since this isn't
coreai-opt's own code. Bumps the trailing line from "Copyright (c) 2023
Apple Inc." (when coremltools first vendored it) to "Copyright (c) 2026
Apple Inc." (when coreai-opt vendored it from coremltools).

Also excludes src/coreai_opt/deps/ from the add-license-header
pre-commit hook, which otherwise unconditionally stamps coreai-opt's
own BSD-3 header onto any .py file lacking one.
Attributes the vendored kmeans1d source under src/coreai_opt/deps/ to
its upstream MIT license, per Apple legal's guidance for incorporating
vendored dependencies (mirrors the equivalent NOTICE.txt commit in
coremltools, which vendored the same code).
The vendored-source-matches-upstream canary hashes our _core.cpp
against coremltools' live upstream copy. Now that coreai-opt's copy
carries its own vendoring-year copyright line ("Copyright (c) 2026
Apple Inc.") rather than coremltools' ("Copyright (c) 2023 Apple
Inc."), the raw byte hashes never match. Strip that one known,
intentional line from both sides before hashing so the canary still
catches real drift in the kmeans1d logic without permanently failing
on the copyright year.
The unanchored `deps/` rule (from the initial commit, alongside other
personal-scratch entries like scratch/ and /local_stash) silently
shadowed the newly-added, tracked src/coreai_opt/deps/ package: search
tools like ripgrep skip gitignored paths by default, and any future
new file added under deps/ would need an explicit force-add. Nothing
in the repo (Makefile, CI, scripts, docs) reads or writes a deps/
folder, so remove the rule outright rather than just anchoring it.
The vendored kmeans1d is JIT-compiled via torch.utils.cpp_extension,
which needs a C++ toolchain on the host at runtime; previously this
requirement rode in transitively via the coremltools wheel, so it's
worth calling out now that coremltools is optional.
Drop "(copied byte-for-byte from coremltools)" -- attribution already
lives in the vendored files' headers and NOTICE.txt, and the phrase
will go stale once the core is further optimized (per aseemw's PR apple#31
review).
aseemw's review: comments explaining implementation choices by diffing
against what coremltools' build does are unnecessary now that this is
a fresh adaptation, not a diff against coremltools. Drop the comment
entirely rather than rephrasing it to be self-contained; the flags
themselves are unchanged.
Re-verified via explain-coreai-opt (aseemw asked why docs needs the
coreml group): Sphinx's autodoc/autosummary machinery does genuinely
import every public coreai_opt module at build time, but nothing under
src/coreai_opt/ imports coremltools anywhere -- the only remaining
kmeans1d/coremltools attribution now lives in comments and the vendored
package's own LICENSE/header, none of which autodoc touches. No CI job
builds docs either, so this has no build-time effect to verify against.
aseemw's PR apple#31 review: kmeans1d correctness tests importing
coremltools as a live oracle defeats the purpose of removing it as a
required dependency, and the "must stay byte-identical to upstream"
canary asserts an invariant that isn't actually required (the vendored
code can be modified/adapted going forward).

Collapses TestEquivalence/TestKmeans1DBehavior/TestsKmeans1dFromCoremltools
into one TestKmeans1D class and deletes TestVendoredSourceMatchesUpstream
entirely. Small discrete-input cases (k>n clamping, fewer-than-k-cluster
collapse, explicit duplicates) now assert hardcoded literal values
directly, captured by running the vendored cluster() once (already
validated against coremltools in CI, so a known-good golden value).

Larger/random-array cases (unweighted/weighted equivalence grids,
large-n up to 100k, fp16 duplicates, weighted-integer-counts) would
require thousands of literal values to hardcode inline, so their
inputs and expected coremltools-oracle output are instead checked into
tests/palettization/assets/kmeans1d/*.npz and loaded at test time
(~3.3MB total, mostly the large-n case).

Net effect: this test file has zero coremltools imports (confirmed via
AST inspection), so coreml stops being an implicit test dependency for
kmeans1d correctness.
@guru-desh
guru-desh force-pushed the u/gurudesh/remove-coremltools-dependency-torch-cpp branch from 416db85 to cd2e650 Compare July 15, 2026 18:25
@guru-desh
guru-desh merged commit edd4720 into apple:main Jul 15, 2026
8 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

enhancement New feature or request

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants